#coding:utf-8
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from paddleocr import PaddleOCR


def drawRectBox(image, box, label, conf, fontC):
    """
    修改后的函数，同时显示车牌号和置信度
    """
    x1, y1, x2, y2 = box
    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2, cv2.LINE_AA)

    # 将车牌号和置信度组合成显示文本
    display_text = f"{label} ({conf:.2f})"

    # 将OpenCV图像转为PIL图像
    pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(pil_img)

    # 获取文本尺寸
    left, top, right, bottom = fontC.getbbox(display_text)
    text_width = right - left
    text_height = bottom - top

    # 文字背景位置
    text_bg_x1 = x1
    text_bg_y1 = max(0, y1 - text_height - 5)

    # 绘制半透明背景
    text_bg = Image.new('RGBA', (text_width, text_height), (0, 0, 255, 150))
    pil_img.paste(text_bg, (text_bg_x1, text_bg_y1), text_bg)

    # 绘制文字
    draw.text((x1, text_bg_y1), display_text, font=fontC, fill=(255, 255, 255))

    return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

def img_cvread(img_path):
    """读取图片（支持中文路径）"""
    # 使用 cv2.imdecode 读取，避免中文路径报错
    img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
    return img

def get_license_result(ocr,image):
    """
    image:输入的车牌截取照片
    输出，车牌号与置信度
    """
    result = ocr.ocr(image, cls=True)[0]
    if result:
        license_name, conf = result[0][1]
        if '·' in license_name:
            license_name = license_name.replace('·', '')
        return license_name, conf
    else:
        return None, None

# 需要检测的图片地址
img_path = "../testimage/2.jpg"
now_img = img_cvread(img_path)

fontC = ImageFont.truetype("../Font/platech.ttf", 50, 0)
# 加载ocr模型
cls_model_dir = 'paddleModels/whl/cls/ch_ppocr_mobile_v2.0_cls_infer'
rec_model_dir = 'paddleModels/whl/rec/ch/ch_PP-OCRv4_rec_infer'
ocr = PaddleOCR(use_angle_cls=False, lang="ch", det=False, cls_model_dir=cls_model_dir,rec_model_dir=rec_model_dir)

# 所需加载的模型目录
path = 'yolov8sbest.pt'
# 加载预训练模型
# conf	0.25	object confidence threshold for detection
# iou	0.7	int.ersection over union (IoU) threshold for NMS
model = YOLO(path, task='detect')
# model = YOLO(path, task='detect',conf=0.5)
# 检测图片
results = model(img_path)[0]

location_list = results.boxes.xyxy.tolist()
if len(location_list) >= 1:
    location_list = [list(map(int, e)) for e in location_list]
    # 截取每个车牌区域的照片
    license_imgs = []
    for each in location_list:
        x1, y1, x2, y2 = each
        cropImg = now_img[y1:y2, x1:x2]
        license_imgs.append(cropImg)
    # 车牌识别结果
    lisence_res = []
    conf_list = []
    for each in license_imgs:
        license_num, conf = get_license_result(ocr, each)
        if license_num:
            lisence_res.append(license_num)
            conf_list.append(conf)
        else:
            lisence_res.append('无法识别')
            conf_list.append(0)
    for (text, conf), box in zip(zip(lisence_res, conf_list), location_list):
        now_img = drawRectBox(now_img, box, text, conf, fontC)

now_img = cv2.resize(now_img,dsize=None,fx=0.5,fy=0.5,interpolation=cv2.INTER_LINEAR)
cv2.imshow("YOLOv8 Detection", now_img)
cv2.waitKey(0)
